#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import os
import re
import sys


class ReplaceText:
    def __init__(self, old, new) -> None:
        self.ptn = re.compile(old)
        self.new = new

    def generate_new_path(self, path):
        name = os.path.basename(path)  # /a/b/123.txt -> 123.txt
        root, ext = os.path.splitext(name)  # 123.txt -> 123 .txt
        new_name = self.ptn.sub(self.new, root) + ext  #
        return os.path.join(os.path.dirname(path), new_name)

    def rename(self, path, is_debug: bool = False):
        # 跳过无意义的路径
        if ".@__thumb" in path:
            return
        if "DVD_MENU" in path:
            return

        new_path = self.generate_new_path(path)
        if path == new_path:
            return

        if ".sc." in new_path:
            new_path = new_path.replace(".sc.", ".zh.")
        elif ".SC." in new_path:
            new_path = new_path.replace(".SC.", ".zh.")

        if is_debug:
            print("[Debug]", path, "->", new_path)
            return

        print("[*]", path, "->", new_path, end="")
        if os.path.exists(new_path):
            print(" - 已存在")
            return

        os.rename(path, new_path)
        print(" - OK")


def rename(args):
    path = args.path
    target = args.target
    new_str = args.new
    is_debug = args.debug

    if not os.path.exists(path):
        print(path, "文件不存在")
        return

    renamer = ReplaceText(target, new_str)
    if os.path.isfile(path):
        renamer.rename(path, is_debug)
        return

    for root, _, files in os.walk(path):
        for f in files:
            renamer.rename(os.path.join(root, f), is_debug)


video_ext_list = [".mp4", ".mkv", ".avi", ".flv", ".wmv", ".mov"]
sub_ext_list = [".srt", ".sub", ".ass", ".ssa", ".idx", ".smi", ".usf"]


def is_video_file(path):
    return os.path.splitext(path)[1].lower() in video_ext_list


def is_sub_file(path):
    return os.path.splitext(path)[1].lower() in sub_ext_list


def auto_index(args):
    """
    1. 自动编号
    2. 根据季，增加前缀S00E，
    3. 自动处理字幕，如果是 ass 开头
    """
    path = args.path
    season = args.season
    is_debug = args.debug

    if not os.path.exists(path):
        print(path, "文件不存在")
        return

    video_list = []
    ass_list = []
    for root, _, files in os.walk(path):
        for f in files:
            path = os.path.join(root, f)
            if ".@__thumb" in path:
                continue
            filename = f.lower()

            if filename.endswith(".ass"):
                ass_list.append(path)
            elif is_video_file(path):
                video_list.append(path)

    video_list.sort()
    ass_list.sort()

    gen_new_names(video_list, season, is_debug)
    gen_new_names(ass_list, season, is_debug)


def gen_new_names(video_list, season, is_debug=True):
    size = len(video_list)
    if season < 10:
        season_str = "S0" + str(season) + "E"
    else:
        season_str = "S" + str(season) + "E"

    for i, old_path in enumerate(video_list):
        root = os.path.dirname(old_path)

        ext = os.path.splitext(old_path)[1]
        if ext.lower() == ".ass":
            ext = ".zh.ass"

        # 根据文件的数量生成编号，如果小于100，则前面补0，如 01；如果大于100，则前面补0，如 001
        i = i + 1
        name = str(i)
        if size < 100:  # 01~99
            if i < 10:  # 00~09
                name = "0" + name
        elif size < 1000:  # 001 ~ 999
            if i < 10:  # 000~009
                name = "00" + name
            elif i < 100:  # 010~099
                name = "0" + name
        new_path = os.path.join(root, season_str + name + ext)

        is_ok = name in old_path and name in new_path

        if is_debug:
            print("[DEBUG]", old_path, "->", new_path, is_ok)
            continue

        print("[*]", old_path, "->", new_path, end="")
        if os.path.exists(new_path):
            print(new_path, " - 已存在")
            continue

        print(" - OK")
        os.rename(old_path, new_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NAS Rename Tool")
    subparsers = parser.add_subparsers(help="子命令")

    parser_rename = subparsers.add_parser(
        "rename",
        help="重命名文件",
    )
    parser_auto = subparsers.add_parser(
        "auto",
        help="自动编号，从 01 开始",
    )

    # 重命名
    parser_rename.add_argument("path", type=str, help="目录")
    parser_rename.add_argument(
        "-t",
        "--target",
        help="目标字符串，支持正则匹配。注意，正则需要双引号。",
    )
    parser_rename.add_argument(
        "-n",
        "--new",
        default="",
        required=False,
        help="新字符串，使用新的字符串替换目标字符串。默认为空字符串。",
    )
    parser_rename.add_argument(
        "-d",
        "--debug",
        action="store_true",
        help="测试模式，不会修改文件名。",
    )

    # 自动编号
    parser_auto.add_argument("path", type=str, help="目录")
    parser_auto.add_argument(
        "-s",
        "--season",
        type=int,
        default=1,
        required=False,
        help="季数，默认为 1",
    )
    parser_auto.add_argument(
        "-d",
        "--debug",
        action="store_true",
        help="测试模式，不会修改文件名。",
    )

    if len(sys.argv) <= 1:
        parser.print_help()
        exit(1)

    args = parser.parse_args()

    cmd = sys.argv[1]
    if cmd == "auto":
        if len(sys.argv) == 2:
            parser_auto.print_help()
            exit(1)
        auto_index(args)
    elif cmd == "rename":
        rename(args)
    else:
        print("未知命令", cmd)
